import pyrealsense2 as rs
from tqdm import tqdm

def print_supported_formats():
    ctx = rs.context()
    devices = ctx.query_devices()
    for i, dev in enumerate(devices):
        print(f"Device {i}: {dev}")
    dev = devices[0]
    for sensor in dev.sensors:
        print(f"sensor: {sensor}")
        for profile in sensor.get_stream_profiles():
            vid_profile = profile.as_video_stream_profile()
            stream_type = profile.stream_type()
            assert stream_type in [rs.stream.infrared, rs.stream.depth, rs.stream.color]  # D435上仅有这三种数据流, 红外, 深度, rgb
            name = stream_type.name
            width = vid_profile.width()
            height = vid_profile.height()
            fps = vid_profile.fps()
            fmt = profile.format().name
            if fmt not in ['bgr8', 'z16']: continue
            # print(f"    {vid_profile}")
            print(f"    {name}, {width}x{height}, {fps=}, {fmt}")

def find_compatible_depth_profiles(target_rgb_profile: dict):
    """ 通过rgb配置找到对应可用的depth配置
    Args:
        target_rgb_profile: 包含'width', 'height', 'fps'关键字的dict
    """
    pipeline = rs.pipeline()
    config = rs.config()
    ctx = rs.context()
    if not ctx.query_devices():
        print("[WARNING] No available RealSense devices")
        return
    dev = ctx.query_devices()[0]
    all_depth_profiles = []
    for profile in dev.first_depth_sensor().get_stream_profiles():
        vid_profile = profile.as_video_stream_profile()
        if profile.format().name != 'z16': continue
        all_depth_profiles.append({
            'width': vid_profile.width(),
            'height': vid_profile.height(),
            'fps': vid_profile.fps(),
            'format': profile.format().name
        })
    compatible_depth_profiles = []
    for dp in tqdm(all_depth_profiles):
        test_config = rs.config()
        test_config.enable_stream(rs.stream.color, target_rgb_profile['width'], target_rgb_profile['height'], rs.format.bgr8, target_rgb_profile['fps'])
        test_config.enable_stream(rs.stream.depth, dp['width'], dp['height'], rs.format.z16, dp['fps'])
        try:
            test_config.resolve(pipeline)
            compatible_depth_profiles.append(dp)
        except RuntimeError:
            pass
    print(f"{len(all_depth_profiles)=}")
    print(f"{len(compatible_depth_profiles)=}")
    for dp in compatible_depth_profiles:
        print(f"{dp['width']:>4}x{dp['height']:<4}, fps: {dp['fps']:>3}")


if __name__ == '__main__':
    print_supported_formats()
    rgb_profile = {
        'width': 1920,
        'height': 1080,
        'fps': 30,
    }
    # find_compatible_depth_profiles(rgb_profile)

"""
USB3.0下, 用lsusb -t -v查看D435的速度为5000Mbps
1920x1080 30fps支持所有depth配置

1280x720 , fps:  30
1280x720 , fps:  15
1280x720 , fps:   6
 848x480 , fps:  90
 848x480 , fps:  60
 848x480 , fps:  30
 848x480 , fps:  15
 848x480 , fps:   6
 848x100 , fps: 300
 848x100 , fps: 100
 640x480 , fps:  90
 640x480 , fps:  60
 640x480 , fps:  30
 640x480 , fps:  15
 640x480 , fps:   6
 640x360 , fps:  90
 640x360 , fps:  60
 640x360 , fps:  30
 640x360 , fps:  15
 640x360 , fps:   6
 480x270 , fps:  90
 480x270 , fps:  60
 480x270 , fps:  30
 480x270 , fps:  15
 480x270 , fps:   6
 424x240 , fps:  90
 424x240 , fps:  60
 424x240 , fps:  30
 424x240 , fps:  15
 424x240 , fps:   6
 256x144 , fps: 300
 256x144 , fps:  90

全部可用配置:
sensor: <pyrealsense2.sensor: "Stereo Module">
    depth, 1280x720, fps=30, z16
    depth, 1280x720, fps=15, z16
    depth, 1280x720, fps=6, z16
    depth, 848x480, fps=90, z16
    depth, 848x480, fps=60, z16
    depth, 848x480, fps=30, z16
    depth, 848x480, fps=15, z16
    depth, 848x480, fps=6, z16
    depth, 848x100, fps=300, z16
    depth, 848x100, fps=100, z16
    depth, 640x480, fps=90, z16
    depth, 640x480, fps=60, z16
    depth, 640x480, fps=30, z16
    depth, 640x480, fps=15, z16
    depth, 640x480, fps=6, z16
    depth, 640x360, fps=90, z16
    depth, 640x360, fps=60, z16
    depth, 640x360, fps=30, z16
    depth, 640x360, fps=15, z16
    depth, 640x360, fps=6, z16
    depth, 480x270, fps=90, z16
    depth, 480x270, fps=60, z16
    depth, 480x270, fps=30, z16
    depth, 480x270, fps=15, z16
    depth, 480x270, fps=6, z16
    depth, 424x240, fps=90, z16
    depth, 424x240, fps=60, z16
    depth, 424x240, fps=30, z16
    depth, 424x240, fps=15, z16
    depth, 424x240, fps=6, z16
    depth, 256x144, fps=300, z16
    depth, 256x144, fps=90, z16
sensor: <pyrealsense2.sensor: "RGB Camera">
    color, 1920x1080, fps=30, bgr8
    color, 1920x1080, fps=15, bgr8
    color, 1920x1080, fps=6, bgr8
    color, 1280x720, fps=30, bgr8
    color, 1280x720, fps=15, bgr8
    color, 1280x720, fps=6, bgr8
    color, 960x540, fps=60, bgr8
    color, 960x540, fps=30, bgr8
    color, 960x540, fps=15, bgr8
    color, 960x540, fps=6, bgr8
    color, 848x480, fps=60, bgr8
    color, 848x480, fps=30, bgr8
    color, 848x480, fps=15, bgr8
    color, 848x480, fps=6, bgr8
    color, 640x480, fps=60, bgr8
    color, 640x480, fps=30, bgr8
    color, 640x480, fps=15, bgr8
    color, 640x480, fps=6, bgr8
    color, 640x360, fps=60, bgr8
    color, 640x360, fps=30, bgr8
    color, 640x360, fps=15, bgr8
    color, 640x360, fps=6, bgr8
    color, 424x240, fps=60, bgr8
    color, 424x240, fps=30, bgr8
    color, 424x240, fps=15, bgr8
    color, 424x240, fps=6, bgr8
    color, 320x240, fps=60, bgr8
    color, 320x240, fps=30, bgr8
    color, 320x240, fps=6, bgr8
    color, 320x180, fps=60, bgr8
    color, 320x180, fps=30, bgr8
    color, 320x180, fps=6, bgr8
"""
